#!/usr/bin/env python3
"""
Script to merge DOCX files from multiple directories into a single DOCX file.
Only processes files starting with 'AA' and excludes subdirectories.
"""

import os
import sys
import argparse
import zipfile
import tempfile
import shutil
import xml.etree.ElementTree as ET
from docx import Document
from docx.shared import Pt
from docx.enum.text import WD_PARAGRAPH_ALIGNMENT
from docx.enum.dml import MSO_THEME_COLOR_INDEX
import warnings


def extract_text_from_xml(filepath):
    """
    Extract text directly from DOCX XML when python-docx fails.
    
    Args:
        filepath (str): Path to the DOCX file
        
    Returns:
        list: List of text paragraphs, or None if extraction failed
    """
    try:
        temp_dir = tempfile.mkdtemp()
        paragraphs = []
        
        # Extract the DOCX file
        with zipfile.ZipFile(filepath, 'r') as zip_ref:
            zip_ref.extractall(temp_dir)
        
        # Read the document.xml file
        doc_xml_path = os.path.join(temp_dir, 'word', 'document.xml')
        if not os.path.exists(doc_xml_path):
            shutil.rmtree(temp_dir)
            return None
        
        # Parse the XML
        tree = ET.parse(doc_xml_path)
        root = tree.getroot()
        
        # Define namespaces
        namespaces = {
            'w': 'http://schemas.openxmlformats.org/wordprocessingml/2006/main'
        }
        
        # Find all paragraph elements (w:p)
        paragraph_elements = root.findall('.//w:p', namespaces)
        
        for para_elem in paragraph_elements:
            # Find all text elements within this paragraph
            text_elements = para_elem.findall('.//w:t', namespaces)
            
            # Combine all text in this paragraph
            paragraph_text = ""
            for text_elem in text_elements:
                if text_elem.text:
                    paragraph_text += text_elem.text
            
            # Add paragraph if it has content
            if paragraph_text.strip():
                paragraphs.append(paragraph_text.strip())
        
        # Clean up
        shutil.rmtree(temp_dir)
        
        return paragraphs if paragraphs else None
        
    except Exception as e:
        print(f"    XML extraction failed: {str(e)}")
        if 'temp_dir' in locals():
            shutil.rmtree(temp_dir, ignore_errors=True)
        return None


def repair_docx_file(filepath):
    """
    Attempt to repair a corrupted DOCX file by extracting and re-packaging.
    
    Args:
        filepath (str): Path to the corrupted DOCX file
        
    Returns:
        str: Path to repaired file, or None if repair failed
    """
    try:
        # Create a temporary directory
        temp_dir = tempfile.mkdtemp()
        
        # Extract the DOCX file (it's a zip archive)
        with zipfile.ZipFile(filepath, 'r') as zip_ref:
            zip_ref.extractall(temp_dir)
        
        # Check if the extracted files contain the required DOCX structure
        required_files = ['word/document.xml', 'word/styles.xml', '[Content_Types].xml']
        missing_files = []
        
        for req_file in required_files:
            if not os.path.exists(os.path.join(temp_dir, req_file)):
                missing_files.append(req_file)
        
        if missing_files:
            print(f"    Missing required files: {', '.join(missing_files)}")
            shutil.rmtree(temp_dir)
            return None
        
        # Create a new repaired DOCX file
        repaired_path = filepath + '.repaired.docx'
        
        # Create a new zip file with the extracted contents
        with zipfile.ZipFile(repaired_path, 'w', zipfile.ZIP_DEFLATED) as zip_out:
            for root, dirs, files in os.walk(temp_dir):
                for file in files:
                    file_path = os.path.join(root, file)
                    arcname = os.path.relpath(file_path, temp_dir)
                    zip_out.write(file_path, arcname)
        
        # Clean up
        shutil.rmtree(temp_dir)
        
        return repaired_path
        
    except Exception as e:
        print(f"    Repair attempt failed: {str(e)}")
        if 'temp_dir' in locals():
            shutil.rmtree(temp_dir, ignore_errors=True)
        return None


def get_aa_docx_files(directories):
    """
    Get all DOCX files starting with 'AA' from specified directories.
    Excludes subdirectories and temporary files.
    
    Args:
        directories (list): List of directory paths to search
        
    Returns:
        list: List of tuples (filepath, filename, directory)
    """
    all_files = []
    
    for directory in directories:
        if not os.path.isdir(directory):
            print(f"Warning: Directory '{directory}' does not exist, skipping...")
            continue
        
        print(f"Scanning directory: {os.path.abspath(directory)}")
        
        # Get all files in the directory (not subdirectories)
        try:
            files = os.listdir(directory)
        except PermissionError:
            print(f"Warning: Cannot access directory '{directory}', skipping...")
            continue
        
        aa_files = []
        for f in files:
            filepath = os.path.join(directory, f)
            
            # Only process files (not directories)
            if os.path.isfile(filepath):
                # Check if file starts with 'AA' and ends with '.docx'
                if f.startswith('AA') and f.lower().endswith('.docx'):
                    # Skip temporary files
                    if not f.startswith('.~') and '~' not in f:
                        aa_files.append((filepath, f, directory))
        
        aa_files.sort()  # Sort files alphabetically
        all_files.extend(aa_files)
        
        if aa_files:
            print(f"  Found {len(aa_files)} files starting with 'AA':")
            for i, (_, filename, _) in enumerate(aa_files, 1):
                print(f"    {i}. {filename}")
        else:
            print(f"  No files starting with 'AA' found in this directory")
    
    return all_files


def merge_docx_files(files_list, output_file):
    """
    Merge DOCX files from files list into a single output file.
    Preserves formatting and skips corrupted files.
    
    Args:
        files_list (list): List of tuples (filepath, filename, directory)
        output_file (str): Path to output merged DOCX file
    """
    if not files_list:
        print("No files to merge!")
        return
    
    # Suppress warnings during processing
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        merged_doc = Document()
    
    print(f"\nFound {len(files_list)} files to merge:")
    for i, (_, filename, directory) in enumerate(files_list, 1):
        print(f"  {i}. {filename} (from: {directory})")
    
    successful_merges = 0
    
    # Merge each document
    for i, (filepath, filename, directory) in enumerate(files_list, 1):
        print(f"\nProcessing {filename} ({i}/{len(files_list)})...")
        print(f"  From: {directory}")
        
        # Try to open the file normally first
        src_doc = None
        repaired_file = None
        original_filepath = filepath
        
        try:
            # Test if file can be opened
            test_doc = Document(filepath)
            print(f"  ✓ File is readable, proceeding with merge...")
            
            # Open the source document again for actual processing
            src_doc = Document(filepath)
        except Exception as e:
            error_msg = str(e)
            if "NULL" in error_msg or "archive" in error_msg:
                print(f"  ✗ Error processing {filename}: Corrupted DOCX file structure")
                print(f"    Attempting to repair...")
                
                # Try to repair the file
                repaired_file = repair_docx_file(filepath)
                if repaired_file:
                    print(f"  ✓ File repaired successfully, using repaired version")
                    try:
                        src_doc = Document(repaired_file)
                        filepath = repaired_file  # Use repaired file for processing
                    except Exception as repair_error:
                        print(f"  ✗ Repaired file still cannot be opened: {str(repair_error)}")
                        src_doc = None
                
                # If repair failed or repaired file still can't be opened, try XML extraction
                if src_doc is None:
                    print(f"  ✗ Repair failed, trying XML text extraction...")
                    
                    # Try to extract text directly from XML
                    text_paragraphs = extract_text_from_xml(original_filepath)
                    if text_paragraphs:
                        print(f"  ✓ XML extraction successful, creating text-only version")
                        # Create a simple document with the extracted text
                        src_doc = Document()
                        for para_text in text_paragraphs:
                            if para_text.strip():
                                src_doc.add_paragraph(para_text)
                    else:
                        print(f"  ✗ XML extraction failed, skipping this file")
                        src_doc = None
            else:
                print(f"  ✗ Error processing {filename}: {error_msg}")
                src_doc = None
        
        if src_doc is None:
            print(f"    Skipping this file and continuing...")
            # Clean up repaired file if it exists
            if repaired_file and os.path.exists(repaired_file):
                os.remove(repaired_file)
            continue
        
        try:
            # Suppress warnings during processing
            with warnings.catch_warnings():
                warnings.simplefilter("ignore")
                
                # Add page break before each document except the first
                if i > 1:
                    merged_doc.add_page_break()
                
                # Add filename as heading with directory info
                heading = merged_doc.add_paragraph()
                heading_run = heading.add_run(f"{filename}")
                heading_run.bold = True
                heading_run.font.size = Pt(16)
                heading.alignment = WD_PARAGRAPH_ALIGNMENT.CENTER
                
                # Add directory as subheading
                dir_para = merged_doc.add_paragraph()
                dir_run = dir_para.add_run(f"Source: {directory}")
                dir_run.italic = True
                dir_run.font.size = Pt(10)
                dir_para.alignment = WD_PARAGRAPH_ALIGNMENT.CENTER
                
                # Add space after heading
                merged_doc.add_paragraph()
                
                # Copy paragraphs with enhanced formatting preservation
                para_count = 0
                for paragraph in src_doc.paragraphs:
                    para_count += 1
                    if para_count % 1000 == 0:
                        print(f"    Processing paragraph {para_count}/{len(src_doc.paragraphs)}...")
                    
                    if paragraph.text.strip() == '':
                        merged_doc.add_paragraph()
                        continue
                    
                    new_para = merged_doc.add_paragraph()
                    
                    # Copy paragraph alignment and basic formatting
                    new_para.alignment = paragraph.alignment
                    
                    # Copy paragraph spacing if available
                    if paragraph.paragraph_format:
                        pf = paragraph.paragraph_format
                        tf = new_para.paragraph_format
                        
                        if hasattr(pf, 'line_spacing') and pf.line_spacing:
                            tf.line_spacing = pf.line_spacing
                        if hasattr(pf, 'space_before') and pf.space_before:
                            tf.space_before = pf.space_before
                        if hasattr(pf, 'space_after') and pf.space_after:
                            tf.space_after = pf.space_after
                    
                    # Copy runs with formatting
                    for run in paragraph.runs:
                        if run.text:
                            new_run = new_para.add_run(run.text)
                            new_run.bold = run.bold
                            new_run.italic = run.italic
                            new_run.underline = run.underline
                            
                            # Copy font formatting with error handling
                            try:
                                if run.font.name:
                                    new_run.font.name = run.font.name
                                if run.font.size:
                                    new_run.font.size = run.font.size
                                if run.font.color and hasattr(run.font.color, 'rgb'):
                                    new_run.font.color.rgb = run.font.color.rgb
                                if hasattr(run.font, 'highlight_color') and run.font.highlight_color:
                                    new_run.font.highlight_color = run.font.highlight_color
                            except Exception as font_error:
                                # Skip font formatting if it causes errors
                                pass
                    
                    # Preserve heading styles
                    if paragraph.style and paragraph.style.name:
                        style_name = paragraph.style.name.lower()
                        if style_name.startswith('heading'):
                            try:
                                # Apply heading formatting
                                for run in new_para.runs:
                                    if not run.bold:  # Don't override existing bold
                                        run.bold = True
                                    if hasattr(run.font, 'size') and not run.font.size:
                                        run.font.size = Pt(14)
                            except:
                                pass
                
                # Copy tables
                for table in src_doc.tables:
                    try:
                        new_table = merged_doc.add_table(rows=len(table.rows), cols=len(table.columns))
                        
                        # Copy table alignment
                        if hasattr(table, 'alignment'):
                            new_table.alignment = table.alignment
                        
                        for i, row in enumerate(table.rows):
                            for j, cell in enumerate(row.cells):
                                target_cell = new_table.cell(i, j)
                                target_cell.text = ""  # Clear default text
                                
                                # Copy cell content with formatting
                                for para in cell.paragraphs:
                                    if para.text.strip():
                                        cell_para = target_cell.add_paragraph()
                                        cell_para.alignment = para.alignment
                                        
                                        for run in para.runs:
                                            if run.text:
                                                cell_run = cell_para.add_run(run.text)
                                                cell_run.bold = run.bold
                                                cell_run.italic = run.italic
                                                cell_run.underline = run.underline
                                                
                                                try:
                                                    if run.font.name:
                                                        cell_run.font.name = run.font.name
                                                    if run.font.size:
                                                        cell_run.font.size = run.font.size
                                                except:
                                                    # Skip font formatting if it causes errors
                                                    pass
                    except Exception as table_error:
                        print(f"    Warning: Error copying table: {str(table_error)}")
                        continue
            
            print(f"  ✓ Successfully merged {filename} ({len(src_doc.paragraphs)} paragraphs, {len(src_doc.tables)} tables)")
            successful_merges += 1
            
        except Exception as e:
            error_msg = str(e)
            if "WD_COLOR_INDEX" in error_msg or "XML mapping" in error_msg:
                print(f"  ✗ Error processing {filename}: Color formatting issue")
                print(f"    This file has unsupported color formatting that cannot be merged")
            elif "NULL" in error_msg or "archive" in error_msg:
                print(f"  ✗ Error processing {filename}: Corrupted DOCX file structure")
                print(f"    This file may be damaged or have formatting issues")
            else:
                print(f"  ✗ Error processing {filename}: {error_msg}")
            print(f"    Skipping this file and continuing...")
            continue
        finally:
            # Clean up repaired file if it exists
            if repaired_file and os.path.exists(repaired_file):
                os.remove(repaired_file)
    
    # Save the merged document
    try:
        merged_doc.save(output_file)
        print(f"\n✓ Successfully created merged document: {output_file}")
        print(f"  Files successfully merged: {successful_merges}/{len(files_list)}")
        if successful_merges < len(files_list):
            print(f"  Note: {len(files_list) - successful_merges} files were skipped due to errors")
    except Exception as e:
        print(f"  ✗ Error saving merged document: {str(e)}")


def main():
    parser = argparse.ArgumentParser(
        description='Merge DOCX files starting with "AA" from multiple directories into a single DOCX file',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog='''
Examples:
  python3 merge_docx.py dir1 dir2 dir3 -o merged.docx
  python3 merge_docx.py /path/to/docs /path/to/more/docs -o output.docx
        '''
    )
    parser.add_argument('directories', nargs='+', 
                       help='Directories to search for AA*.docx files')
    parser.add_argument('-o', '--output', default='merged_document.docx', 
                       help='Output filename (default: merged_document.docx)')
    
    args = parser.parse_args()
    
    # Validate directories
    valid_dirs = []
    for directory in args.directories:
        if os.path.isdir(directory):
            valid_dirs.append(directory)
        else:
            print(f"Warning: Directory '{directory}' does not exist and will be skipped")
    
    if not valid_dirs:
        print("Error: No valid directories specified")
        sys.exit(1)
    
    output_file = os.path.abspath(args.output)
    
    print("DOCX File Merger - AA Files Only")
    print("=" * 50)
    print(f"Searching for AA*.docx files in:")
    for i, directory in enumerate(valid_dirs, 1):
        print(f"  {i}. {os.path.abspath(directory)}")
    print(f"Output file: {output_file}")
    print("-" * 50)
    
    # Get all AA docx files
    files_list = get_aa_docx_files(valid_dirs)
    
    if not files_list:
        print("No files starting with 'AA' found in any of the specified directories")
        sys.exit(1)
    
    # Merge the files
    merge_docx_files(files_list, output_file)


if __name__ == "__main__":
    main()